import gym
import numpy as np
from scipy.linalg import solve_continuous_are
import time
# 定义系统参数
g = 9.81  # 重力加速度
m = 0.1   # 杆的质量
M = 1.0   # 小车的质量
l = 0.5   # 杆的长度

# 状态向量 x = [x, dx, theta, dtheta]
# 控制输入 u = F (小车的水平推力)

# 线性化系统的状态矩阵 A 和输入矩阵 B
A = np.array([
    [0, 1, 0, 0],
    [0, 0, -0.9, 0],
    [0, 0, 0, 1],
    [0, 0, 20.64, 0]
])

B = np.array([
    [0],
    [0.9091],
    [0],
    [2.064]
])
print(A)
print(B)
# 权重矩阵 Q 和 R
Q = np.diag([1, 1, 100, 1])  # 调整权重，重点关注角度
R = np.array([[0.01]])

# 求解黎卡提方程 (Riccati equation)
P = solve_continuous_are(A, B, Q, R)
K = np.linalg.inv(R) @ B.T @ P

# 初始化环境
env = gym.make("CartPole-v1", render_mode="human")
state, info = env.reset()

# 控制器运行
try:
    for t in range(100000):
        # 计算控制输入
        u = -K @ state

        # 将控制输入转换为动作 (0: 向左, 1: 向右)
        action = 0 if u > 0 else 1

        # 执行动作
        state, reward, done, truncated, info = env.step(action)

        # 渲染环境
        env.render()

        # 检查是否完成
        if done:
            print(f"Episode finished after {t + 1} timesteps")
            break

        # 添加一个小的延迟以便观察
        time.sleep(0.01)
except KeyboardInterrupt:
    print("Simulation stopped manually after {} timesteps".format(t))

env.close()